import os
import pickle
import random

import numpy as np
from PIL import Image
import torch
from torch.utils.data import Dataset
from torchvision import transforms
import torch.nn.functional as F


def apply_trans(input):
    trans = transforms.Compose([
        transforms.RandomHorizontalFlip(0.5),
        transforms.RandomAffine(degrees=0, translate=(0.1,0.1), scale=(0.75,1.25), resample=Image.NEAREST, fillcolor=None)
    ])

    return trans(input)


class NameAndImage(Dataset):
    def __init__(self, data_dir, sample=[], length=128):
        self.data_dir = data_dir

        self.image_dir = os.path.join(data_dir, f'images')

        self.filenames = os.listdir(self.image_dir)
        self.filenames.sort()

        self.filenames_ = []
        for s in sample:
            self.filenames_ += self.filenames[s*length: (s+1)*length]
        self.filenames = self.filenames_

    def __len__(self):
        return len(self.filenames)

    def __getitem__(self, index):
        filename = self.filenames[index]

        image = Image.open(f'{self.image_dir}/{filename}').convert('L')
        # [C,H,W] and dynamic range [0,255], W & H must match G output resolution
        image = np.array(image)[None, :]
        image = torch.tensor(image, dtype=torch.float32)

        return filename[:-4], image


class ISIC2018Task1_w_mask_(Dataset):
    def __init__(self, save_dir, which_set='training'):
        assert which_set in ['training', 'validation']

        self.save_dir = save_dir
        self.image_dir = f'./data/images/{which_set}'
        self.mask_dir = f'./data/masks/{which_set}'
        self.filenames = os.listdir(f'./data/images/{which_set}')
        self.filenames.sort()
        self.filenames = [name.split('.')[0] for name in self.filenames]
        self.len = 2594 if which_set=='training' else 100

    def __len__(self):
        return self.len

    def __getitem__(self, index):
        filename = self.filenames[index]

        with open(f"{os.path.join(self.save_dir, 'ws', filename)}.pkl", 'rb') as f:
            w = pickle.load(f)['w']

        image = Image.open(f'{self.image_dir}/{filename}.png').convert('RGB')
        image = np.array(image).transpose([2, 0, 1])/255
        image = torch.tensor(image, dtype=torch.float32)

        mask = Image.open(f'{self.mask_dir}/{filename}_segmentation.png').convert('L')
        mask = np.array(mask)/255
        mask = torch.tensor(mask, dtype=torch.float32).unsqueeze(0)  # H, W -> C, H, W

        return w, image, mask


class ISIC2018Task1_w_mask(Dataset):
    def __init__(self, save_dir):
        self.save_dir = save_dir
        self.image_dir = f'./data/images/'
        self.mask_dir = f'./data/masks/'

        training_set = os.listdir(f'{self.image_dir}/training')
        training_set.sort()
        validation_set = os.listdir(f'{self.image_dir}/validation')
        validation_set.sort()
        self.filenames = training_set + validation_set    # 2594 + 100
        self.filenames = [name.split('.')[0] for name in self.filenames]

    def __len__(self):
        return 2694

    def __getitem__(self, index):
        which_set = 'training' if index<2594 else 'validation'

        filename = self.filenames[index]

        w = torch.zeros((1))
        if self.save_dir:
            with open(f"{os.path.join(self.save_dir, 'ws', filename)}.pkl", 'rb') as f:
                w = pickle.load(f)['w']

        image = Image.open(f'{self.image_dir}/{which_set}/{filename}.png').convert('RGB')
        image = np.array(image).transpose([2, 0, 1])/255
        image = torch.tensor(image, dtype=torch.float32)

        mask = Image.open(f'{self.mask_dir}/{which_set}/{filename}_segmentation.png').convert('L')
        mask = np.array(mask)/255
        mask = torch.tensor(mask, dtype=torch.float32).unsqueeze(0)  # H, W -> C, H, W

        return w, image, mask


class Handmask_(Dataset):
    def __init__(self, save_dir):
        self.save_dir = save_dir
        self.image_dir = os.path.join(save_dir, 'handmark', f'images')
        self.mask_dir = os.path.join(save_dir, 'handmark', f'masks')
        self.repre_dir = os.path.join(save_dir, 'handmark', f'repres')

        self.filenames = os.listdir(self.image_dir)
        self.filenames.sort(key=lambda x:int(x[:-4]))

    def __len__(self):
        return len(self.filenames)

    def __getitem__(self, index):
        filename = self.filenames[index]

        with open(os.path.join(self.repre_dir, f"{filename.split('.')[0]}.pkl"), 'rb') as f:
            repre = pickle.load(f)['repre']
        for i in range(len(repre)):
            repre[i] = transforms.Resize((256, 256), Image.BILINEAR)(repre[i])
        repre = torch.cat(repre, dim=1).squeeze(0)

        mask = Image.open(os.path.join(self.mask_dir, filename)).convert('L')
        mask = np.array(mask)/255
        mask = torch.tensor(mask, dtype=torch.float32).unsqueeze(0)  # H, W -> C, H, W

        return repre, mask


class Handmask(Dataset):
    def __init__(self, handmark_dir, multi_class=False, combine=False):
        self.multi_class = multi_class
        self.combine = combine
        assert multi_class and combine

        self.image_dir = os.path.join(handmark_dir, f'images')
        self.mask_dir = os.path.join(handmark_dir, f'masks')
        self.ws_dir = os.path.join(handmark_dir, f'ws')

        self.filenames = os.listdir(self.image_dir)
        self.filenames.sort(key=lambda x:int(x[:-4]))

    def __len__(self):
        return len(self.filenames)

    def __getitem__(self, index):
        filename = self.filenames[index]

        with open(os.path.join(self.ws_dir, f'{filename[:-4]}.pkl'), 'rb') as f:
            w = pickle.load(f)['w']

        # image = Image.open(os.path.join(self.image_dir, filename)).convert('RGB')
        # image = np.array(image).transpose([2, 0, 1])/25
        image = Image.open(f'{self.image_dir}/{filename}').convert('L')
        image = np.array(image)[None, :]/255
        image = torch.tensor(image, dtype=torch.float32)

        mask = np.array(Image.open(os.path.join(self.mask_dir, filename)).convert('L'))
        mask = torch.tensor(mask, dtype=torch.long)

        # H, W = mask.shape
        # if self.multi_class:
        #     m = torch.zeros(H, W)
        #     for label in zip([2,3,41,42], [1,2,3,4]):
        #         m = torch.tensor(np.where(mask==label[0], label[1], m), dtype=torch.long)
        #     mask = m

        #     if self.combine:
        #         mask = torch.tensor(np.where(mask==3, 1, mask))
        #         mask = torch.tensor(np.where(mask==4, 2, mask))
        # else:
        #     mask = torch.tensor(mask/255, dtype=torch.float32).unsqueeze(0)  # H, W -> 1, H, W

        return w, image, mask


class Handmask_2d(Dataset):
    def __init__(self, handmark_dir, multi_class=False, combine=False):
        self.multi_class = multi_class
        self.combine = combine

        self.image_dir = os.path.join(handmark_dir, f'images')
        self.mask_dir = os.path.join(handmark_dir, f'masks')
        self.ws_dir = os.path.join(handmark_dir, f'ws')

        self.filenames = os.listdir(self.image_dir)
        self.filenames.sort(key=lambda x:int(x[:-4]))

    def __len__(self):
        return len(self.filenames)

    def __getitem__(self, index):
        filename = self.filenames[index]

        with open(os.path.join(self.ws_dir, f'{filename[:-4]}.pkl'), 'rb') as f:
            w = pickle.load(f)['w']

        # image = Image.open(os.path.join(self.image_dir, filename)).convert('RGB')
        # image = np.array(image).transpose([2, 0, 1])/25
        image = Image.open(f'{self.image_dir}/{filename}').convert('L')
        image = np.array(image)[None, :]/255
        image = torch.tensor(image, dtype=torch.float32)

        mask = np.array(Image.open(os.path.join(self.mask_dir, filename)).convert('L'))[None, :]/255
        mask = torch.tensor(mask, dtype=torch.float32)

        return w, image, mask

class Supervised(Dataset):
    def __init__(self, data_dir, multi_class):
        self.multi_class = multi_class

        self.image_dir = os.path.join(data_dir, 'images')
        self.mask_dir = os.path.join(data_dir, 'masks')

        self.filenames = os.listdir(self.image_dir)

    def __len__(self):
        return len(self.filenames)

    def __getitem__(self, index):
        filename = self.filenames[index]

        image = Image.open(f'{self.image_dir}/{filename}').convert('RGB')
        image = np.array(image).transpose([2, 0, 1])/255
        image = torch.tensor(image, dtype=torch.float32)

        mask = np.array(Image.open(os.path.join(self.mask_dir, filename)).convert('L'))
        if self.multi_class:
            masks = []
            for label in [2,41,3,42]:
                masks.append(torch.tensor(np.where(mask==label, 1.0, 0.0), dtype=torch.float32).unsqueeze(0))  # H, W -> 1, H, W
            mask = torch.cat(masks, dim=0)  # 1, H, W -> C, H, W
        else:
            mask = torch.tensor(mask/255, dtype=torch.float32).unsqueeze(0)  # H, W -> 1, H, W

        return image, mask


class Triplet(Dataset):
    def __init__(self, data_dir, multi_class, aug, sample=[], save_dir=None, return_w=False, combine=False, length=64, w_steps=1000):
        self.data_dir = data_dir
        self.aug = aug
        
        assert multi_class and combine

        if save_dir and return_w:
            self.w_dir = os.path.join(save_dir, f'ws_{w_steps}')
            assert os.path.exists(self.w_dir)
        else:
            self.w_dir = None

        self.image_dir = os.path.join(data_dir, f'images')
        self.mask_dir = os.path.join(data_dir, f'masks')

        self.filenames = os.listdir(self.image_dir)
        self.filenames.sort()

        self.filenames_ = []
        for s in sample:
            self.filenames_ += self.filenames[s*length: (s+1)*length]
        self.filenames = self.filenames_

    def __len__(self):
        return len(self.filenames)

    def __getitem__(self, index):
        filename = self.filenames[index]

        w = torch.zeros((1))
        if self.w_dir:
            with open(f"{os.path.join(self.w_dir, filename.split('.')[0])}.pkl", 'rb') as f:
                w = pickle.load(f)['w']

        # image = Image.open(f'{self.image_dir}/{filename}').convert('RGB')
        # image = np.array(image).transpose([2, 0, 1])/255
        image = Image.open(f'{self.image_dir}/{filename}').convert('L')
        image = np.array(image)[None, :]/255
        image = torch.tensor(image, dtype=torch.float32)

        mask = Image.open(os.path.join(self.mask_dir, filename)).convert('L')
        mask = torch.tensor(np.array(mask), dtype=torch.long)
        H, W = mask.shape
        C = image.shape[0]

        if self.aug:
            mask = mask.unsqueeze(0)
            image, mask = apply_trans(torch.cat([image, mask], dim=0)).split([C, 1], dim=0)
            mask = mask.squeeze(0).to(torch.long)

        return w, image, mask


class US(Dataset):
    def __init__(self, data_dir, aug, sample=[]):
        self.data_dir = data_dir
        self.aug = aug

        self.image_dir = os.path.join(data_dir, f'images')
        self.mask_dir = os.path.join(data_dir, f'masks')

        self.filenames = os.listdir(self.image_dir)
        self.filenames.sort(key=lambda x:int(x[:-4]))
        
        self.filenames_ = [self.filenames[s] for s in sample]
        self.filenames = self.filenames_
    
    def __len__(self):
        return len(self.filenames)

    def __getitem__(self, index):
        filename = self.filenames[index]

        image = Image.open(f'{self.image_dir}/{filename}').convert('L')
        image = np.array(image)[None, :]/255
        image = torch.tensor(image, dtype=torch.float32)

        mask = Image.open(os.path.join(self.mask_dir, filename)).convert('L')
        mask = torch.tensor(np.array(mask)[None, :]/255, dtype=torch.float32)

        C = image.shape[0]
        if self.aug:
            image, mask = apply_trans(torch.cat([image, mask], dim=0)).split([C, 1], dim=0)

        return image, mask


class MR_slice(Dataset):
    def __init__(self, data_dir, aug, sample, length=128):
        self.data_dir = data_dir
        self.aug = aug

        self.image_dir = os.path.join(data_dir, f'images')
        self.mask_dir = os.path.join(data_dir, f'masks')

        self.filenames = os.listdir(self.image_dir)
        self.filenames.sort()
        print(len(self.filenames))

        filenames_ = []
        for s in sample:
            filenames_ += self.filenames[s*length: (s+1)*length]
        self.filenames = filenames_

    def __len__(self):
        return len(self.filenames)

    def __getitem__(self, index):
        filename = self.filenames[index]

        image = Image.open(os.path.join(self.image_dir, filename)).convert('L')
        image = np.array(image)[None, :]/255
        image = torch.tensor(image, dtype=torch.float32)

        mask = Image.open(os.path.join(self.mask_dir, filename)).convert('L')
        mask = torch.tensor(np.array(mask)[None, :], dtype=torch.long)

        C = image.shape[0]
        if self.aug:
            image, mask = apply_trans(torch.cat([image, mask], dim=0)).split([C, 1], dim=0)
            mask = mask.to(torch.long)
        mask = mask.squeeze(0)

        return image, mask


class Ds(Dataset):
    def __init__(self, data_dir, sub, aug):
        self.data_dir = data_dir
        self.aug = aug
        self.sub = sub
        assert sub=='ss' or sub=='sd'

        self.image_dirs = [os.path.join(data_dir, f'{sub}_{domain}') for domain in range(3)]
        self.mask_dir = os.path.join(data_dir, f'masks')

        self.filenames = os.listdir(os.path.join(self.data_dir, 'ss_0'))
        self.filenames.sort()
        self.base_len = len(self.filenames)
    def __len__(self):
        return self.base_len*3

    def __getitem__(self, index):
        domain = index // self.base_len
        n = index % self.base_len
        filename = self.filenames[n]

        image = Image.open(os.path.join(self.image_dirs[domain], filename)).convert('L')
        image = np.array(image)[None, :]/255
        image = torch.tensor(image, dtype=torch.float32)

        mask = Image.open(os.path.join(self.mask_dir, filename)).convert('L')
        mask = torch.tensor(np.array(mask)[None, :], dtype=torch.long)

        C = image.shape[0]
        if self.aug:
            image, mask = apply_trans(torch.cat([image, mask], dim=0)).split([C, 1], dim=0)
            mask = mask.squeeze(0).to(torch.long)

        return image, mask


if __name__ == "__main__":
    # t = Triplet(
    #     'data/CANDI-128-256',
    #     True,
    #     True,
    #     [1],
    #     '',
    #     False,

    #     True,
    #     resize=True
    # )
    # test = t.__getitem__(0)
    # print(test[1].shape)
    # print(test[2].shape)
    # from torchvision.utils import save_image
    # save_image(test[1], 'test1.png')
    # save_image(test[2]*1.0, 'test2.png')
    t = MR_slice(
        data_dir='data/CANDI-128-160',
        aug=False,
        sample=[103]
    )
    print(t.filenames)
    print(t.__len__())


CANDI_label = [
            'Cerebral_WM', # 灰质
            'Cerebral_CX', # 白质
            'Lateral-Vent',
            'Cerebellum-WM',
            'Cerebellum-CX',
            'Thalamus-Proper',
            'Caudate',
            'Putamen',
            'Pallidum',
            '3rd-Vent',
            '4rd-Vent',
            'Brain-Stem',
            'Hippocampus',
            'Amygdala',
            'CSF',
            'VentralDC'
        ]